{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Control-Lyapunov Function for Driving a Wheeled Robot"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This notebook implements the controller proposed in \"[Closed loop steering of unicycle-like vehicles via Lyapunov techniques](https://ieeexplore.ieee.org/abstract/document/388294)\" by M. Aicardi, G. Casalino, A. Bicchi, and A. Balestrino."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook Setup \n",
    "The following cell will install Drake, checkout the underactuated repository, and set up the path (only if necessary).\n",
    "- On Google's Colaboratory, this **will take approximately two minutes** on the first time it runs (to provision the machine), but should only need to reinstall once every 12 hours.  Colab will ask you to \"Reset all runtimes\"; say no to save yourself the reinstall.\n",
    "- On Binder, the machines should already be provisioned by the time you can run this; it should return (almost) instantly.\n",
    "\n",
    "More details are available [here](http://underactuated.mit.edu/drake.html)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "try:\n",
    "    import pydrake\n",
    "    import underactuated\n",
    "except ImportError:\n",
    "    !curl -s https://raw.githubusercontent.com/RussTedrake/underactuated/master/scripts/setup/jupyter_setup.py > jupyter_setup.py\n",
    "    from jupyter_setup import setup_underactuated\n",
    "    setup_underactuated()\n",
    "\n",
    "# Setup matplotlib.  \n",
    "from IPython import get_ipython\n",
    "if get_ipython() is not None: get_ipython().run_line_magic(\"matplotlib\", \"inline\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# python libraries\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.patches as patches\n",
    "\n",
    "# pydrake imports\n",
    "from pydrake.all import (VectorSystem, DiagramBuilder, SymbolicVectorSystem, LogOutput,\n",
    "                         Variable, Simulator, cos, sin)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Robot Kinematics\n",
    "We start writing down symbolically the kinematics of the robot.\n",
    "The vehicle will be then represented as a Drake `SymbolicVectorSystem`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# state of the robot (in cartesian coordinates)\n",
    "z1 = Variable(\"z1\") # horizontal position\n",
    "z2 = Variable(\"z2\") # vertical position\n",
    "z3 = Variable(\"z3\") # angular position\n",
    "cartesian_state = [z1, z2, z3]\n",
    "\n",
    "# control input of the robot\n",
    "u1 = Variable(\"u1\") # linear velocity\n",
    "u2 = Variable(\"u2\") # angular velocity\n",
    "input = [u1, u2]\n",
    "\n",
    "# nonlinear dynamics, the whole state is measured (output = state)\n",
    "dynamics = [u1*cos(z3), u1*sin(z3), u2]\n",
    "robot = SymbolicVectorSystem(\n",
    "    state=cartesian_state,\n",
    "    input=input,\n",
    "    output=cartesian_state,\n",
    "    dynamics=dynamics,\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Write the Control Law\n",
    "It is time to write the controller.\n",
    "To do that we construct a `VectorSystem` with a method called `DoCalcVectorOutput`.\n",
    "In the latter, we write all the steps needed to compute the instantaneous control action from the vehicle state."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# you will write the control law in it in the next cell\n",
    "# by defining the function called \"lyapunov_controller\"\n",
    "class ParkingController(VectorSystem):\n",
    "\n",
    "    def __init__(self, lyapunov_controller):\n",
    "        # 3 inputs (robot state)\n",
    "        # 2 outputs (robot inputs)\n",
    "        VectorSystem.__init__(self, 3, 2)\n",
    "        self.lyapunov_controller = lyapunov_controller\n",
    "\n",
    "    def DoCalcVectorOutput(\n",
    "            self,\n",
    "            context,          # not used\n",
    "            cartesian_state,  # input of the controller\n",
    "            controller_state, # not used\n",
    "            input             # output of the controller\n",
    "        ):\n",
    "        \n",
    "        # upack state of the robot\n",
    "        z1, z2, z3 = cartesian_state\n",
    "        \n",
    "        # state in polar coordinates\n",
    "        x1 = np.sqrt(z1**2 + z2**2) # radial coordinate\n",
    "        x2 = np.arctan2(z2, z1)     # angular coordinate\n",
    "        x3 = x2 - z3\n",
    "        \n",
    "        # evaluate the function below and return the robot's input\n",
    "        input[:] = self.lyapunov_controller(x1, x2, x3)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In the following cell, modify the function `lyapunov_controller` to implement the control law from point (b) of the exercise.\n",
    "The function must return the linear $u_1$ and the angular $u_2$ velocity of the robot.\n",
    "Currently, it just returns zero input, and the robot does not move.\n",
    "\n",
    "To simplify your work, the input of this function is the state already expressed in polar coordinates:\n",
    "- `x1`: radial coordinate,\n",
    "- `x2`: angular coordinate,\n",
    "- `x3`: angle between the radial vector and the robot orientation.\n",
    "\n",
    "**Attention:**\n",
    "The function $u_2 = \\pi_2 (\\mathbf x)$ is not defined for $x_3 = 0$.\n",
    "To prevent divisions by zero, we defined a tolerance of $10^{-3}$ (do not modify this value) and whenever the absolute value of $x_3$ is lower than or equal to this tolerance you should use the Taylor approximation $\\cos(x_3)\\sin(x_3)/x_3 \\approx 1$."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def lyapunov_controller(x1, x2, x3):\n",
    "    if np.abs(x3) <= 1e-3: # do not modify\n",
    "        u1 = 0 # modify here\n",
    "        u2 = 0 # modify here\n",
    "    else:\n",
    "        u1 = 0 # modify here\n",
    "        u2 = 0 # modify here\n",
    "    return u1, u2"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Wire the Block Diagram\n",
    "Now we connect the controller with the robot.\n",
    "To this end, we construct a Drake diagram."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# construction site for our closed-loop system\n",
    "builder = DiagramBuilder()\n",
    "\n",
    "# add the robot to the diagram\n",
    "# the method .AddSystem() simply returns a pointer to the system\n",
    "# we passed as input, so it's ok to give it the same name\n",
    "robot = builder.AddSystem(robot)\n",
    "\n",
    "# add the controller\n",
    "controller = builder.AddSystem(ParkingController(lyapunov_controller))\n",
    "\n",
    "# wire the controller with the system\n",
    "builder.Connect(robot.get_output_port(0), controller.get_input_port(0))\n",
    "builder.Connect(controller.get_output_port(0), robot.get_input_port(0))\n",
    "\n",
    "# add a logger to the diagram\n",
    "# this will store the state trajectory\n",
    "logger = LogOutput(robot.get_output_port(0), builder)\n",
    "\n",
    "# complete the construction of the diagram\n",
    "diagram = builder.Build()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Simulate the Closed-Loop System\n",
    "It's time to simulate the closed-loop system: just pass the diagram we constructed to the simulator and press play!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# reset logger to delete old trajectories\n",
    "# this is needed if you want to simulate the system multiple times\n",
    "logger.reset()\n",
    "\n",
    "# set up a simulation environment\n",
    "simulator = Simulator(diagram)\n",
    "\n",
    "# set the initial cartesian state to a random initial position\n",
    "# try initial_state = np.random.randn(3) for a random initial state\n",
    "initial_state = [0, 1, 0] \n",
    "context = simulator.get_mutable_context()\n",
    "context.SetContinuousState(initial_state)\n",
    "\n",
    "# simulate from zero to sim_time\n",
    "# the trajectory will be stored in the logger\n",
    "sim_time = 20.\n",
    "simulator.AdvanceTo(sim_time)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Plot the Parking Trajectory\n",
    "No need to fully inderstand the following lines of code, they are just plotting helpers...\n",
    "Just get a sense of which parameter to modify in case you want to costumize the plot."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# vehicle size in the plot\n",
    "length = .5\n",
    "width = .3\n",
    "\n",
    "# function that draws a triangle given the cartesian state of the robot\n",
    "def draw_robot(cartesian_state, **kwargs):\n",
    "    \n",
    "    # transformation matrix from global to robot frame\n",
    "    z1, z2, z3 = cartesian_state\n",
    "    T = np.array([\n",
    "        [np.cos(z3), -np.sin(z3), z1],\n",
    "        [np.sin(z3), np.cos(z3), z2]\n",
    "    ])\n",
    "\n",
    "    # compute cartesian coordinates of the corners of the car\n",
    "    left_back = T.dot([-length/2, width/2, 1])\n",
    "    right_back = T.dot([-length/2, -width/2, 1])\n",
    "    front = T.dot([length/2, 0, 1])\n",
    "    \n",
    "    # draw a triangle that symbolizes the robot\n",
    "    robot = patches.Polygon(\n",
    "        np.vstack([left_back, right_back, front]),\n",
    "        facecolor='none',\n",
    "        **kwargs\n",
    "    )\n",
    "    plt.gca().add_patch(robot)\n",
    "\n",
    "# function that draws the trajectory of the robot\n",
    "# initial and final states are red and blue, respectively\n",
    "def draw_trajectory(trajectory):\n",
    "\n",
    "    # draw trajectory of the robot in black\n",
    "    for state in trajectory.T:\n",
    "        draw_robot(state, edgecolor='k')\n",
    "        \n",
    "    # draw initial configuration of the robot in red\n",
    "    draw_robot(\n",
    "        trajectory[:,0],\n",
    "        edgecolor='r',\n",
    "        linewidth=3,\n",
    "        label='Initial configuration'\n",
    "    )\n",
    "\n",
    "    # draw final configuration of the robot in blue\n",
    "    draw_robot(\n",
    "        trajectory[:,-1],\n",
    "        edgecolor='b',\n",
    "        linewidth=3,\n",
    "        label='Final configuration'\n",
    "    )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Here is the final plot!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# draw parking trajectory\n",
    "plt.figure(figsize=(10, 10))\n",
    "draw_trajectory(logger.data())\n",
    "\n",
    "# adapt figure limits to the trajectory\n",
    "robot_size = max(length, width)\n",
    "plt.xlim([\n",
    "        min(logger.data()[0,:]) - robot_size,\n",
    "        max(logger.data()[0,:]) + robot_size\n",
    "])\n",
    "plt.ylim([\n",
    "        min(logger.data()[1,:]) - robot_size,\n",
    "        max(logger.data()[1,:]) + robot_size\n",
    "])\n",
    "\n",
    "# misc plot settings\n",
    "plt.gca().set_aspect('equal')\n",
    "plt.legend()\n",
    "plt.xlabel(r'$z_1$')\n",
    "plt.ylabel(r'$z_2$')\n",
    "plt.title('Robot parking trajectory')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Autograding\n",
    "\n",
    "You can check your work by running the following cell:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from underactuated.exercises.lyapunov.control.test_control import TestControl\n",
    "from underactuated.exercises.grader import Grader\n",
    "Grader.grade_output([TestControl], [locals()], 'results.json')\n",
    "Grader.print_test_results('results.json')"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}